import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
from yolo import YOLO
from tqdm import tqdm
from PIL import Image
import xml.etree.ElementTree as ET
from utils.utils import get_classes

if __name__ == '__main__':
    classes_path    = 'model_data/work_classes.txt'
    data_path  = 'D:/work/work_data'
    image_ids = open(os.path.join(data_path, "D:/work/work_data/Image_sets/test.txt")).read().strip().split('\n')
    yolo = YOLO(confidence=0.001, nms_iou=0.5)
    class_names, _ = get_classes(classes_path)
    all_ = 0
    corrects = 0
    for image_id in tqdm(image_ids):
        image_path = os.path.join(data_path, "Images/" + image_id + ".png")
        anno_path = os.path.join(data_path, "annotations/"+ image_id + ".xml")
        root = ET.parse(anno_path).getroot()
        labels = []
        for obj in root.findall('object'):
            obj_name = obj.find('name').text
            labels.append(obj_name)
        image = Image.open(image_path)
        num_correct = yolo.get_general_accuracy(image, labels, class_names)
        corrects += num_correct
        all_ += 5
    print('accuracy: %.2f %%' % (100*corrects/all_))
